#include <mcp/transport/stdio_transport.h>
#include <mcp/error.h>
#include <mcp/json/value.h>

#include <iostream>
#include <sstream>
#include <chrono>
#include <memory>
#include <thread>
#include <atomic>

namespace mcp {
namespace transport {

// StdioTransport implementation
struct StdioTransport::Impl {
    std::atomic<bool> running{false};
    std::atomic<bool> connected{false};
    std::unique_ptr<std::thread> receiver_thread;
    MessageHandler message_handler;
    ErrorHandler error_handler;
    CloseHandler close_handler;
    OpenHandler open_handler;
    size_t buffer_size = 4096;
    std::atomic<size_t> bytes_read{0};
    std::atomic<size_t> bytes_written{0};
    bool use_stderr = false;
    mutable std::mutex transport_mutex;

    Impl(bool use_stderr_) : use_stderr(use_stderr_) {}

    void start_reading();
    void stop_reading();
    void process_input();
};

StdioTransport::StdioTransport() : impl_(std::make_unique<Impl>(false)) {}

StdioTransport::StdioTransport(bool use_stderr) : impl_(std::make_unique<Impl>(use_stderr)) {}

StdioTransport::~StdioTransport() {
    stop();
}

bool StdioTransport::start() {
    std::lock_guard<std::mutex> lock(impl_->transport_mutex);

    if (impl_->running) {
        return true;
    }

    impl_->running = true;
    impl_->connected = true;

    // Start the receiver thread
    impl_->receiver_thread = std::make_unique<std::thread>([this]() {
        impl_->start_reading();
    });

    // Notify open handler
    if (impl_->open_handler) {
        impl_->open_handler();
    }

    return true;
}

void StdioTransport::stop() {
    std::lock_guard<std::mutex> lock(impl_->transport_mutex);

    if (!impl_->running) {
        return;
    }

    impl_->running = false;
    impl_->connected = false;

    // Stop the receiver thread
    impl_->stop_reading();

    if (impl_->receiver_thread && impl_->receiver_thread->joinable()) {
        impl_->receiver_thread->join();
        impl_->receiver_thread.reset();
    }

    // Notify close handler
    if (impl_->close_handler) {
        impl_->close_handler();
    }
}

bool StdioTransport::is_connected() const {
    return impl_->connected;
}

bool StdioTransport::send(const json::Value& message) {
    if (!impl_->running || !impl_->connected) {
        return false;
    }

    try {
        std::string serialized = message.serialize();

        // Write to stdout or stderr
        std::ostream& output_stream = impl_->use_stderr ? std::cerr : std::cout;
        output_stream << serialized << std::endl;

        if (output_stream.fail()) {
            if (impl_->error_handler) {
                impl_->error_handler(Error::io_error("Failed to write to stdio"));
            }
            return false;
        }

        impl_->bytes_written += serialized.size();
        return true;
    } catch (const std::exception& e) {
        if (impl_->error_handler) {
            impl_->error_handler(Error::serialization_error(e.what()));
        }
        return false;
    }
}

std::coroutine_handle<> StdioTransport::send_async(const json::Value& message) {
    // For stdio transport, send_async behaves the same as send
    // In a real implementation, this would use coroutines for async I/O
    send(message);
    return std::noop_coroutine();
}

void StdioTransport::set_message_handler(MessageHandler handler) {
    impl_->message_handler = std::move(handler);
}

void StdioTransport::set_error_handler(ErrorHandler handler) {
    impl_->error_handler = std::move(handler);
}

void StdioTransport::set_close_handler(CloseHandler handler) {
    impl_->close_handler = std::move(handler);
}

void StdioTransport::set_open_handler(OpenHandler handler) {
    impl_->open_handler = std::move(handler);
}

std::string StdioTransport::get_type() const {
    return "stdio";
}

json::Value StdioTransport::get_capabilities() const {
    json::Value capabilities = json::Value::object();
    capabilities["type"] = "stdio";
    capabilities["supports_async"] = false;
    capabilities["supports_compression"] = false;
    capabilities["buffer_size"] = static_cast<int64_t>(impl_->buffer_size);
    return capabilities;
}

void StdioTransport::set_buffer_size(size_t size) {
    impl_->buffer_size = size;
}

size_t StdioTransport::get_bytes_read() const {
    return impl_->bytes_read;
}

size_t StdioTransport::get_bytes_written() const {
    return impl_->bytes_written;
}

void StdioTransport::Impl::start_reading() {
    while (running) {
        try {
            process_input();
        } catch (const std::exception& e) {
            if (error_handler) {
                error_handler(Error::io_error(e.what()));
            }
            break;
        }
    }
}

void StdioTransport::Impl::stop_reading() {
    // No specific action needed for stdio
    // The loop in start_reading will exit when running is false
}

void StdioTransport::Impl::process_input() {
    std::string line;
    std::istream& input_stream = std::cin;

    // Read a line of input
    if (!std::getline(input_stream, line)) {
        // End of input or error
        running = false;
        connected = false;
        if (close_handler) {
            close_handler();
        }
        return;
    }

    // Skip empty lines
    if (line.empty()) {
        return;
    }

    bytes_read += line.size();

    // Parse the JSON message
    try {
        auto result = json::Value::parse(line);
        if (result && message_handler) {
            message_handler(*result);
        } else if (!result && error_handler) {
            error_handler(result.error());
        }
    } catch (const std::exception& e) {
        if (error_handler) {
            error_handler(Error::parse_error(e.what()));
        }
    }
}

// StdioTransportFactory implementation
std::unique_ptr<Transport> StdioTransportFactory::create() {
    return std::make_unique<StdioTransport>();
}

std::string StdioTransportFactory::get_type() const {
    return "stdio";
}

std::vector<std::string> StdioTransportFactory::get_supported_schemes() const {
    return {"stdio", "stdin", "stdout"};
}

// Utility functions for stdio transport
std::unique_ptr<StdioTransport> create_stdio_transport(bool use_stderr) {
    return std::make_unique<StdioTransport>(use_stderr);
}

std::unique_ptr<StdioTransport> create_stdio_transport_with_buffer(size_t buffer_size, bool use_stderr) {
    auto transport = std::make_unique<StdioTransport>(use_stderr);
    transport->set_buffer_size(buffer_size);
    return transport;
}

} // namespace transport
} // namespace mcp